from PIL import Image
import cv2
import numpy as np
import os
import pandas as pd
import docx2txt
import PyPDF2
import openpyxl
from pptx import Presentation

class MultiModalService:
    def __init__(self):
        pass
        
    def process_document(self, file_path):
        """处理各种文档文件"""
        file_ext = os.path.splitext(file_path)[1].lower()
        
        try:
            if file_ext == '.txt':
                with open(file_path, 'r', encoding='utf-8') as f:
                    return f.read()
                    
            elif file_ext == '.docx':
                return docx2txt.process(file_path)
                
            elif file_ext == '.pdf':
                text = ""
                with open(file_path, 'rb') as f:
                    pdf_reader = PyPDF2.PdfReader(f)
                    for page in pdf_reader.pages:
                        text += page.extract_text()
                return text
                
            elif file_ext in ['.xlsx', '.xls']:
                df = pd.read_excel(file_path)
                return df.to_string()
                
            elif file_ext == '.csv':
                df = pd.read_csv(file_path)
                return df.to_string()
                
            elif file_ext in ['.ppt', '.pptx']:
                text = []
                prs = Presentation(file_path)
                for slide in prs.slides:
                    for shape in slide.shapes:
                        if hasattr(shape, "text"):
                            text.append(shape.text)
                return "\n\n".join(text)
                
            else:
                raise ValueError(f"不支持的文件类型: {file_ext}")
                
        except UnicodeDecodeError:
            # 尝试其他编码
            try:
                with open(file_path, 'r', encoding='gbk') as f:
                    return f.read()
            except:
                raise Exception(f"无法识别文件编码: {file_path}")
    
    def process_image(self, image_path):
        """处理图片文件，返回 base64 编码"""
        try:
            import base64
            with open(image_path, 'rb') as image_file:
                base64_image = base64.b64encode(image_file.read()).decode('utf-8')
            return base64_image
        except Exception as e:
            raise Exception(f"处理图片时出错: {str(e)}")
    
    def process_video(self, video_path):
        """处理视频文件，提取关键帧"""
        try:
            video = cv2.VideoCapture(video_path)
            frames = []
            total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
            
            # 每秒提取一帧
            fps = video.get(cv2.CAP_PROP_FPS)
            interval = int(fps)
            
            for i in range(0, total_frames, interval):
                video.set(cv2.CAP_PROP_POS_FRAMES, i)
                ret, frame = video.read()
                if ret:
                    frames.append(frame)
            
            video.release()
            return frames
        except Exception as e:
            raise Exception(f"处理视频时出错: {str(e)}") 